from flask import Flask,request,jsonify
from flask_jwt_extended import (
    JWTManager,create_access_token,create_refresh_token,
    jwt_required,get_jwt_identity
)
from dotenv import load_dotenv
import os
import time

# 加载环境变量
load_dotenv()


app = Flask(__name__)



# 配置JWT
app.config['JWT_SECRET_KEY'] = os.getenv('JWT_SECRET_KEY')

# 从环境变量获取过期时间（秒）
ACCESS_EXPIRE_SECONDS = int(os.getenv('ACCESS_TOKEN_EXPIRES'))
REFRESH_EXPIRES_SECONDS = int(os.getenv('REFRESH_TOKEN_EXPIRES'))

# 初始化 JWT 管理器
jwt = JWTManager(app)

# 模拟用户数据库
users = {
    "user1": {"password": "admin123", "role": "user"},
    "admin": {"password": "admin456", "role": "admin"}
}

# 登录接口：生成双令牌（用 time 计算过期时间）
@app.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')

    user = users.get(username)
    if not user or user['password'] != password:
        return jsonify({'message': '用户名或者密码错误！'}), 401

    # 计算令牌过期时间（当前时间戳+过期秒数）
    current_timestamp = time.time() # 获取当前时间戳
    access_expires = current_timestamp + ACCESS_EXPIRE_SECONDS
    refresh_expires = current_timestamp + REFRESH_EXPIRES_SECONDS

    # 创建令牌时指定过期时间（expires_delta 支持秒数或者时间戳）
    access_token = create_access_token(
        identity = username,
        expires_delta = ACCESS_EXPIRE_SECONDS # 直接传入秒数（flask-jwt-extended 支持）
    )

    refresh_token = create_refresh_token(
        identity = username,
        expires_delta = REFRESH_EXPIRES_SECONDS
    )

    return jsonify({
        'access_token': access_token,
        'refresh_token': refresh_token,
        'token_type': 'bearer',
        # 可选：返回过期时间戳（便于前端判断）
        "access_expires_at": int(access_expires),
        "refresh_expires_at": int(refresh_expires)
    }), 200

# 刷新令牌
@app.route('/refresh',methods=['POST'])
@jwt_required(refresh=True)
def refresh():
    current_user = get_jwt_identity()
    # 生成新的 access_token（过期时间仍用秒数配置）
    new_access_token = create_access_token(
        identity = current_user,
        expires_delta = ACCESS_EXPIRE_SECONDS
    )
    return jsonify({
        'access_token': new_access_token,
    }),200

# 受保护接口示例
@app.route('/protected', methods=['GET'])
@jwt_required()
def protected():
    current_user = get_jwt_identity()
    return jsonify(logged_in_as=current_user), 200


if __name__ == '__main__':
    app.run(debug=True)